import numpy as np
from base_types import TimeSeriesCS
import itertools

def random_walk(size, step_size=1, delta_t=1, distrib='normal', distrib_args=()):
    '''
    Return a TimeSeriesCS object which is the realization of a simple random 
    walk.
    '''
    steps = getattr(np.random, distrib)(*distrib_args, scale=step_size, size=size)
    data = np.add.accumulate(steps)
    return TimeSeriesCS(delta_t=delta_t, data=data)

def ito_iter(t0=0, delta_t=1, a=None, b=None, inner_steps=10):
    '''
    Yield pairs (t, x(t)) where x(t) comes from an Ito's process defined by the
    stochastic differential equation
    
        dx = a(x, t) * dt + b(x, t) * dW_t
        
    Attributes
    ----------
    delta_t : float
        Time step of each point in simulation
    t0 : float
        Initial time of simulation
    a, b: callable or numbers
        Drift and diffusion functions/coefficients for the process
    inner_steps : int
        Number of iterations of the Euler method for integrating the SDE 
        before yielding each new point 
    '''

    raise NotImplementedError

def ito_series(t0=0, tf=None, delta_t=1, size=None, a=None, b=None, inner_steps=10):
    '''
    Return a TimeSeriesCS object representing a realization of a given Ito 
    process.
    
    The user must provide 3 of 4 parameters t0, tf, delta_t and size. If less
    than 3 parameters are given, the series is initialized assuming that t0=0
    and delta_t=1.
    
    See ito_iter for more information on the other parameters of this function.
    '''

    ito_it = ito_iter(t0, delta_t, a, b, inner_steps)
    data = list(x for (_, x) in itertools.islice(ito_it, size))
    return TimeSeriesCS((t0, tf), data)

if __name__ == '__main__':
    import matplotlib.pyplot as plt
    from pylab import *

    #Random walk
    rw = array(list(random_walk(100).items()))
    print rw
    X = rw[:, 0]
    Y = rw[:, 1]
    plt.plot (X, Y)
    plt.show()
    print "\n\nPressione 'enter' para continuar.\n\n"
    raw_input()

    #Ito's process by interation
